import os
from typing import Dict, Optional, Union
import requests
from qwen_agent.tools.base import BaseTool, register_tool

# 注册一个名为'my_amap_weather'的自定义工具
@register_tool('my_amap_weather')
class MyAmapWeather(BaseTool):
    description = '获取对应城市的天气数据'
    # 定义工具的输入参数
    parameters = [{
        'name': 'location',
        'type': 'string',
        'description': '城市/区具体名称，如`北京市海淀区`请描述为`海淀区`',
        'required': True
    }]

    def __init__(self, cfg: Optional[Dict] = None):
        super().__init__(cfg)

        # 高德天气API的URL模板
        self.url = 'https://restapi.amap.com/v3/weather/weatherInfo?city={city}&key={key}'

        # 导入pandas库用于读取城市编码数据
        import pandas as pd
        # 从ModelScope下载城市编码对照表
        self.city_df = pd.read_excel('./resource/AMap_adcode_citycode.xlsx')

        # 获取高德地图API的访问令牌，优先从配置中获取，否则从环境变量中获取
        self.token = self.cfg.get('token', os.environ.get('AMAP_TOKEN', 'xxxxx'))
        # 确保token不为空，否则抛出异常
        assert self.token != '', 'weather api token must be acquired through ' \
            'https://lbs.amap.com/api/webservice/guide/create-project/get-key and set by AMAP_TOKEN'

    def get_city_adcode(self, city_name):
        # 根据城市名称查找对应的区域编码
        filtered_df = self.city_df[self.city_df['中文名'] == city_name]
        if len(filtered_df['adcode'].values) == 0:
            # 如果找不到对应的城市编码，抛出异常
            raise ValueError(f'location {city_name} not found, availables are {self.city_df["中文名"]}')
        else:
            # 返回城市对应的区域编码
            return filtered_df['adcode'].values[0]

    def call(self, params: Union[str, dict], **kwargs) -> str:
        # 验证并格式化输入参数
        params = self._verify_json_format_args(params)

        # 获取位置参数
        location = params['location']
        # 调用高德天气API获取天气数据
        response = requests.get(self.url.format(city=self.get_city_adcode(location), key=self.token))
        data = response.json()
        if data['status'] == '0':
            # API调用失败时抛出异常
            raise RuntimeError(data)
        else:
            # 从返回数据中提取天气和温度信息
            weather = data['lives'][0]['weather']
            temperature = data['lives'][0]['temperature']
            # 返回格式化的天气信息
            return f'{location}的天气是{weather}温度是{temperature}度。'
